import sys
import glob
import os

try:
    sys.path.append(glob.glob('../carla/dist/carla-*%d.%d-%s.egg' % (
        sys.version_info.major,
        sys.version_info.minor,
        'win-amd64' if os.name == 'nt' else 'linux-x86_64'))[0])
except IndexError:
    pass

import carla
import cv2
import numpy as np
import random
import time
import argparse
import threading
import socket
import pygame
from pygame.locals import *  # 用于键盘事件

# 设置UDP客户端
UDP_IP = "192.168.0.2"  # 树莓派的IP地址
UDP_PORT = 5005         # 选择一个端口

def main():
    pygame.init()
    display = pygame.display.set_mode((512, 256))  # 与相机分辨率匹配
    pygame.display.set_caption("CARLA Camera Feed")
    clock = pygame.time.Clock()

    # 在 main 函数开始处
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(("", UDP_PORT))  # 绑定到所有接口的5005端口

    argparser = argparse.ArgumentParser(
        description=__doc__)
    argparser.add_argument(
        '--host',
        metavar='H',
        default='127.0.0.1',
        help='IP of the host server (default: 127.0.0.1)')
    argparser.add_argument(
        '-p', '--port',
        metavar='P',
        default=2000,
        type=int,
        help='TCP port to listen to (default: 2000)')
    argparser.add_argument(
        '-s', '--speed',
        metavar='FACTOR',
        default=1.0,
        type=float,
        help='rate at which the weather changes (default: 1.0)')
    args = argparser.parse_args()

    # 连接到CARLA服务器
    client = carla.Client(args.host, args.port)
    client.set_timeout(10.0)
    world = client.get_world()

    # 获取可用的车辆模型
    blueprints = world.get_blueprint_library().filter('vehicle')
    vehicle_bp = [bp for bp in blueprints if 'tesla.model3' in bp.id][0]  # 选了个特斯拉model3

    # 获取所有可用的生成点（车道位置）
    spawn_points = world.get_map().get_spawn_points()
    
    # 随机选择一个生成点
    spawn_point = random.choice(spawn_points)

    # 在选定位置生成车辆
    vehicle = world.spawn_actor(vehicle_bp, spawn_point)

    # 创建摄像头蓝图
    camera_bp = world.get_blueprint_library().find('sensor.camera.rgb')
    camera_bp.set_attribute('image_size_x', '512')
    camera_bp.set_attribute('image_size_y', '256')
    camera_bp.set_attribute('fov', '110')
    # 添加帧率控制，设置为每秒10帧
    camera_bp.set_attribute('sensor_tick', '0.1')  # 每0.1秒捕获一帧

    # 假设挂载摄像头到车辆的前方
    camera_transform = carla.Transform(carla.Location(x=1.5, z=1.5))  # 偏移量可以根据需要调整

    # 在车辆上附加摄像头
    camera = world.spawn_actor(camera_bp, camera_transform, attach_to=vehicle)

    # 回调函数来处理图像数据
    def process_image(image):
        # 将CARLA图像转换为Pygame surface
        image_array = np.frombuffer(image.raw_data, dtype=np.uint8)
        image_array = np.reshape(image_array, (image.height, image.width, 4))
        image_array = image_array[:, :, :3]  # 去掉alpha通道
        image_surface = pygame.surfarray.make_surface(image_array.swapaxes(0, 1))
        display.blit(image_surface, (0, 0))
        pygame.display.flip()

        # 将图像编码为 JPEG 格式
        _, jpeg_image = cv2.imencode('.jpg', image_array)
        image_data = jpeg_image.tobytes()               # 将JPEG编码的图像转换为字节数组
        
        # 获取 JPEG 编码后的图像大小（单位：字节）
        image_size = len(image_data)
        print(f"发送的图像大小: {image_size} 字节")

        # 设置每个数据包的大小（以1440字节为例）
        MAX_PACKET_SIZE = 1440
        
        # 分割数据并发送
        for i in range(0, len(image_data), MAX_PACKET_SIZE):
            # 获取当前数据包
            packet = image_data[i:i + MAX_PACKET_SIZE]
            
            # 发送数据包
            sock.sendto(packet, (UDP_IP, UDP_PORT))
            #print(f"Sent packet {i // MAX_PACKET_SIZE + 1}")
        

        # 显示图像
        #cv2.imshow("CARLA Camera Feed", image_array)

        # 及时刷新窗口，适当调整等待时间，避免窗口无响应
        #cv2.waitKey(1)  # 等待1ms以便处理键盘事件和窗口刷新

    # 设置摄像头回调
    camera.listen(lambda image: process_image(image))

    # 设置车辆控制
    control = carla.VehicleControl()
    control.throttle = 0.3  # 设置恒定的低速前进
    control.steer = 0.0
    vehicle.apply_control(control)

    # 等待车辆启动
    time.sleep(2.5)
    control.throttle = 0.12  # 设置恒定的低速前进
    vehicle.apply_control(control)

    def receive_commands():
        while True:
            try:
                data, addr = sock.recvfrom(1024)
                command = data.decode('utf-8')
                print(f"Received command: {command}")
                
                # 使用nonlocal声明来访问外部的control变量
                nonlocal control
                if command == "Turn Left":
                    control.steer = -0.5
                elif command == "Turn Right":
                    control.steer = 0.5
                elif command == "Turn Left Slight":
                    control.steer = -0.2
                elif command == "Turn Right Slight":
                    control.steer = 0.2
                elif command == "Keep Straight":
                    control.steer = 0.0
                
                vehicle.apply_control(control)
            except Exception as e:
                print(f"Error receiving data: {e}")
                time.sleep(0.1)

    # 在main函数中创建接收线程
    receive_thread = threading.Thread(target=receive_commands)
    receive_thread.daemon = True  # 设置为守护线程，主程序退出时会自动结束
    receive_thread.start()

    # 主循环
    try:
        while True:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    return
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_LEFT:
                        control.steer = -0.3
                    elif event.key == pygame.K_RIGHT:
                        control.steer = 0.3
                elif event.type == pygame.KEYUP:
                    if event.key in [pygame.K_LEFT, pygame.K_RIGHT]:
                        control.steer = 0.0
                vehicle.apply_control(control)

            time.sleep(0.1)
            clock.tick(30)  # 限制帧率

    except KeyboardInterrupt:
        print("\n退出中...")
    
    finally:
        # 清理过程
        print("Destroying vehicle and camera...")
        camera.stop()
        camera.destroy()
        vehicle.destroy()
        pygame.quit()
        sock.close()
        recv_sock.close()
        print("Cleanup complete.")

if __name__ == '__main__':
    main()
